from mcp.server.fastmcp import FastMCP, Context
from mcp.types import SamplingMessage, TextContent

# A stateful server is required for server-to-client requests like sampling.
mcp = FastMCP(
    name="mcp-sampling-server",
    stateless_http=False,
)


@mcp.tool()
async def create_story(ctx: Context, topic: str) -> str:
    """
    Creates a short story by asking the client to generate it via sampling.

    Args:
        ctx: The MCP Context, used to communicate with the client.
        topic: The topic for the story.

    Returns:
        The story generated by the client's LLM.
    """
    print(f"-> Server: Tool 'create_story' called with topic: '{topic}'")

    try:
        print(f"-> Server: Sending 'sampling/create' request to client...")

        #    The server delegates the "thinking" to the client.
        result = await ctx.session.create_message(
            messages=[
                SamplingMessage(
                    role="user",
                    content=TextContent(type="text", text=f"Write a very short, three-sentence story about: {topic}"),
                )
            ],
            max_tokens=100,
        )

        if result.content.type == "text":
            return result.content.text
        return str(result.content)

    except Exception as e:
        print(f"-> Server: An error occurred during sampling: {e}")
        return f"Error asking client to generate story: {e}"

# Expose the server as a runnable ASGI application
mcp_app = mcp.streamable_http_app()
